from typing import Union, Optional, List
import ee, eemont
import os
from common.logger import logger
from ee.batch import Task


class AssetManager:
    def __init__(self) -> None:
        pass

    @staticmethod
    def asset_exists(asset_id: str) -> bool:
        """Check the existence of the asset"""
        try:
            ee.data.getAsset(asset_id)
        except ee.EEException:
            exists = False
        else:
            exists = True
        return exists

    def get_asset_list(self, parent, include_parent=True):
        """To recursively list all assets within a directory"""
        parent_asset = ee.data.getAsset(parent)
        parent_id = parent_asset["name"]
        parent_type = parent_asset["type"]
        asset_list = []
        child_assets = ee.data.listAssets({"parent": parent_id})["assets"]
        for child_asset in child_assets:
            child_id = child_asset["name"]
            child_type = child_asset["type"]
            if child_type in ["FOLDER", "IMAGE_COLLECTION"]:
                asset_list.extend(self.get_asset_list(child_id))
            else:
                asset_list.append(child_id)
        if include_parent:
            asset_list.append(parent_id)
        return asset_list

    def move_image_collection(self, src_dir, dst_dir):
        if not self.asset_exists(dst_dir):
            ee.data.createAsset({"type": "ImageCollection"}, dst_dir)
        for idx, child_asset in enumerate(
            self.get_asset_list(src_dir, include_parent=False)
        ):
            ee.data.copyAsset(
                child_asset, os.path.join(dst_dir, str(idx)), allowOverwrite=True
            )
        self.recursive_delete_asset(src_dir)

    def recursive_delete_asset(self, src_dir):
        for child_asset in self.get_asset_list(src_dir):
            ee.data.deleteAsset(child_asset)


class TaskManager:
    def __init__(self) -> None:
        pass

    @staticmethod
    def get_active_task():
        active_task_list = [
            task
            for task in Task.list()
            if task.state
            not in ["COMPLETED", "FAILED", "CANCELLED", "CANCEL_REQUESTED"]
        ]
        return active_task_list

    @staticmethod
    def cancel_active_task():
        active_task_list = [
            task
            for task in Task.list()
            if task.state
            not in ["COMPLETED", "FAILED", "CANCELLED", "CANCEL_REQUESTED"]
        ]
        for task in active_task_list:
            Task.cancel(task)


class ImageSmoother:
    def __init__(self, imgCol, areaThreshold=1000, num_cls=2, circleRadius=10) -> None:
        self.num_cls = num_cls
        self.circleRadius = circleRadius
        self.areaThreshold = areaThreshold
        self.imgCol = imgCol

    def majorityFilter(self, img):
        image = ee.Image(img)
        kernel = ee.Kernel.circle(radius=self.circleRadius)
        crop_lst = ee.List([])
        for c_idx, _ in enumerate(range(self.num_cls)):
            new_value = [0] * self.num_cls
            new_value[c_idx] = 1
            cur_crop = image.remap(list(range(1, self.num_cls + 1)), new_value)
            cur_crop_processed = (
                cur_crop.focalMean(kernel=kernel)
                .addBands(ee.Image(c_idx + 1).float())
                .rename(["ratio", "type"])
            )
            crop_lst = crop_lst.add(cur_crop_processed)
        return ee.ImageCollection(crop_lst).qualityMosaic("ratio").select("type")

    def deSpeckler(self):
        def inner_function(img):
            image = ee.Image(img).int8()
            objectId = image.connectedComponents(
                connectedness=ee.Kernel.square(1), maxSize=128
            )
            objectSize = objectId.select("labels").connectedPixelCount(
                maxSize=128, eightConnected=True
            )
            objectArea = objectSize.multiply(ee.Image.pixelArea())
            smallPatches = objectArea.lt(self.areaThreshold)
            smallPatches = smallPatches.selfMask()
            neighborhoodImg = self.majorityFilter(image)
            smoothedImg = image.where(
                smallPatches.eq(1), neighborhoodImg
            ).copyProperties(image, ["Province", "City", "County"])
            smoothedImg = ee.Image(smoothedImg)
            return smoothedImg

        return self.imgCol.map(inner_function)


class GridGenerator:
    def __init__(self, aoi, nx, ny, projection=None) -> None:
        self.aoi = aoi
        self.nx = nx
        self.ny = ny
        self.projection = projection

    def build_from_corner(self, feature, dx, dy):

        coordinates = feature.geometry().coordinates()

        leftX = ee.Number(coordinates.get(0))
        leftY = ee.Number(coordinates.get(1)).subtract(dy)
        rightX = ee.Number(coordinates.get(0)).add(dx)
        rightY = ee.Number(coordinates.get(1))
        corners = ee.List([leftX, leftY, rightX, rightY])

        if self.projection is None:
            geometry = ee.Geometry.Rectangle(corners)
        else:
            geometry = ee.Geometry.Rectangle(corners, proj=self.projection)

        return ee.Feature(geometry)

    def get_box_grid(self):
        coords = ee.List(self.aoi.coordinates())
        minlon = ee.Number(ee.List(ee.List(coords.get(0)).get(0)).get(0))
        maxlon = ee.Number(ee.List(ee.List(coords.get(0)).get(2)).get(0))
        minlat = ee.Number(ee.List(ee.List(coords.get(0)).get(0)).get(1))
        maxlat = ee.Number(ee.List(ee.List(coords.get(0)).get(2)).get(1))

        dx = maxlon.subtract(minlon).abs().divide(self.nx - 1)
        dy = maxlat.subtract(minlat).abs().divide(self.ny - 1)

        lons = ee.List.sequence(minlon, maxlon, None, self.nx).slice(0, -1)
        lats = ee.List.sequence(maxlat, minlat, None, self.ny).slice(0, -1)

        def outter_map(x):
            def inner_map(y):
                return ee.Feature(ee.Geometry.Point([ee.Number(y), ee.Number(x)]))

            return lons.map(inner_map)

        pcoords = lats.map(outter_map)

        ptsgrid = ee.FeatureCollection(pcoords.flatten())
        grid = ptsgrid.map(lambda p: self.build_from_corner(p, dx, dy).set("corner", p))

        return grid


class LandsatGenerator:
    def __init__(
        self,
        start_time: str,
        end_time: str,
        region: Union[ee.Feature, ee.FeatureCollection, ee.Geometry],
        use_mask: bool = True,
        metric: List[str] = [],
        target_bands: Optional[List[str]] = None,
        sensors: Optional[str] = None,
    ) -> None:
        self.start_time = start_time
        self.end_time = end_time
        self.region = region
        self.use_mask = use_mask
        self.target_bands = target_bands
        self.sensors = sensors
        self.metric = metric

    def getLandsat(self):
        if self.target_bands == None:
            self.target_bands = ["BLUE", "GREEN", "RED", "NIR", "SWIR1", "SWIR2"]
        if self.sensors == None:
            self.sensors = dict(l4=True, l5=True, l7=True, l8=True)
        selected_bands = self.target_bands.extend(self.metric)
        collection8 = (
            ee.ImageCollection("LANDSAT/LC08/C02/T1_L2")
            .filterDate(self.start_time, self.end_time)
            .filterBounds(self.region)
            .preprocess()
            .spectralIndices(self.metric)
            .select(selected_bands)
        )
        collection7 = (
            ee.ImageCollection("LANDSAT/LE07/C02/T1_L2")
            .filterDate(self.start_time, self.end_time)
            .filterBounds(self.region)
            .preprocess()
            .spectralIndices(self.metric)
            .select(selected_bands)
        )
        collection5 = (
            ee.ImageCollection("LANDSAT/LT05/C02/T1_L2")
            .filterDate(self.start_time, self.end_time)
            .filterBounds(self.region)
            .preprocess()
            .spectralIndices(self.metric)
            .select(selected_bands)
        )
        col = collection5.merge(collection7).merge(collection8)
        return ee.ImageCollection(col)

    def prepareL4L5(self, image):
        bandList = ["B1", "B2", "B3", "B4", "B5", "B7", "B6"]
        scaling = [10000, 10000, 10000, 10000, 10000, 10000, 1000]
        scaled = ee.Image(image).select(bandList).divide(ee.Image.constant(scaling))
        validQA = [66, 130, 68, 132]
        mask1 = (
            ee.Image(image)
            .select(["pixel_qa"])
            .remap(validQA, ee.List.repeat(1, len(validQA)), 0)
        )
        mask2 = image.select("radsat_qa").eq(0)
        mask3 = image.select(bandList).reduce(ee.Reducer.min()).gt(0)
        mask4 = image.select("sr_atmos_opacity").unmask().lt(300)
        return (
            ee.Image(image)
            .addBands(scaled)
            .updateMask(mask1.And(mask2).And(mask3).And(mask4))
        )

    def prepareL7(image):
        bandList = ["B1", "B2", "B3", "B4", "B5", "B7", "B6"]
        scaling = [10000, 10000, 10000, 10000, 10000, 10000, 1000]
        scaled = ee.Image(image).select(bandList).divide(ee.Image.constant(scaling))
        validQA = [66, 130, 68, 132]
        mask1 = (
            ee.Image(image)
            .select(["pixel_qa"])
            .remap(validQA, ee.List.repeat(1, len(validQA)), 0)
        )
        mask2 = image.select("radsat_qa").eq(0)
        mask3 = image.select(bandList).reduce(ee.Reducer.min()).gt(0)
        mask4 = image.select("sr_atmos_opacity").unmask().lt(300)
        mask5 = ee.Image(image).mask().reduce(ee.Reducer.min()).focal_min(2.5)
        return (
            ee.Image(image)
            .addBands(scaled)
            .updateMask(mask1.And(mask2).And(mask3).And(mask4).And(mask5))
        )

    def prepareL8(image):
        bandList = ["B2", "B3", "B4", "B5", "B6", "B7", "B10"]
        scaling = [10000, 10000, 10000, 10000, 10000, 10000, 1000]
        validTOA = [66, 68, 72, 80, 96, 100, 130, 132, 136, 144, 160, 164]
        validQA = [322, 386, 324, 388, 836, 900]
        scaled = ee.Image(image).select(bandList).divide(ee.Image.constant(scaling))
        mask1 = (
            ee.Image(image)
            .select(["pixel_qa"])
            .remap(validQA, ee.List.repeat(1, len(validQA)), 0)
        )
        mask2 = image.select("radsat_qa").eq(0)
        mask3 = image.select(bandList).reduce(ee.Reducer.min()).gt(0)
        mask4 = (
            ee.Image(image)
            .select(["sr_aerosol"])
            .remap(validTOA, ee.List.repeat(1, len(validQA)), 0)
        )
        return (
            ee.Image(image)
            .addBands(scaled)
            .updateMask(mask1.And(mask2).And(mask3).And(mask4))
        )


def composite_by_month(imgCl, region, num_year, reducer, month_index=None):
    """create composite based on month
    Note that the image should come from the same year!

    Parameters
    ----------
    imgCl: ee.imageCollection,
        the image collection to be composed
    region: ee.Geometry
        the region for the composition
    num_year: int
        the number of years for the composition
    reducer: ee.Reducer
        the way to reduce images within a single month

    Returns
    ----------
    ee.ImageCollection
        a collection containing all composed images
    """
    # determine the initial year
    ini_year = imgCl.first().date().get("year")
    c = ee.List([])
    for y in range(num_year):
        # select images from a single year
        imgCl_subset = imgCl.filterDate(
            ee.Date.fromYMD(ini_year + y, 1, 1), ee.Date.fromYMD(ini_year + y + 1, 1, 1)
        )
        # generate composite for each month
        if month_index == None:
            for i in range(12):
                # select images from a single month
                imgs = imgCl_subset.filterBounds(region).filter(
                    ee.Filter.calendarRange(i, i + 1, "month")
                )
                if imgs.size().getInfo() != 0:
                    # assign the date of the last image within the interval as the 'system:time_start' to the composite
                    date_last = imgs.aggregate_array("system:time_start").reduce(
                        ee.Reducer.last()
                    )
                    print(ee.Date(date_last).format().getInfo())
                    # assign the date appears most to the composite
                    date_mode = ee.Date(
                        imgs.aggregate_array("system:time_start").reduce(
                            ee.Reducer.mode()
                        )
                    ).format()
                    c = c.add(
                        imgs.reduce(reducer)
                        .clip(region)
                        .set({"system:time_start": date_last, "date_mode": date_mode})
                    )
        else:
            imgs = imgCl_subset.filterBounds(region).filter(
                ee.Filter.calendarRange(month_index[0], month_index[1], "month")
            )
            if imgs.size().getInfo() != 0:
                # assign the date of the last image within the interval as the 'system:time_start' to the composite
                date_last = imgs.aggregate_array("system:time_start").reduce(
                    ee.Reducer.last()
                )
                print(ee.Date(date_last).format().getInfo())
                # assign the date appears most to the composite
                date_mode = ee.Date(
                    imgs.aggregate_array("system:time_start").reduce(ee.Reducer.mode())
                ).format()
                c = c.add(
                    imgs.reduce(reducer)
                    .clip(region)
                    .set({"system:time_start": date_last, "date_mode": date_mode})
                )
    return ee.ImageCollection(c)


if __name__ == "__main__":
    import fire

    ee.Initialize()
    fire.Fire(
        {
            "cancel_task": TaskManager.cancel_active_task(),
            "cancel_item": AssetManager().recursive_delete_asset(),
        }
    )
